#!/usr/bin/env python
# -*- coding:utf-8 -*-
import pefile
from datetime import datetime
import re
from sapi import malstr
from sapi import malfunc
from lib.utils import printer

peSection = {
    '.text':"可执行代码",
    '.data':"全局读写数据",
    '.rdata':"只读字符串",
    '.rsrc':"字符串和图标等资源",
    '.reolc':"硬编码地址重定位"
}

def pe_header(path):
    pe = pefile.PE(path)
    printer.info("PE头信息分析: ")
    dMagic = hex(pe.DOS_HEADER.e_magic)
    dlfanew = hex(pe.DOS_HEADER.e_lfanew)
    nSignture = hex(pe.NT_HEADERS.Signature)
    fTimeDateStamp = str(datetime.fromtimestamp(pe.FILE_HEADER.TimeDateStamp))
    fMachine = hex(pe.FILE_HEADER.Machine)
    if dMagic == "0x5a4d":
        printer.test("PE文件的标志位:MZ [十六进制:0x5a4d]")
        printer.test("PE文件头的开始地址(RVA): " + dlfanew)
        if nSignture == "0x4550":
            printer.test("PE文件标识头: PE00 [十六进制:0x4550]")
        else:
            printer.test("PE文件标识头: " + nSignture)
        printer.test("文件创建时间: " + fTimeDateStamp)
        if fMachine == "0x14c":
            printer.test("PE文件：32位x86程序")
        elif fMachine == "0x8664":
            printer.test("PE文件: 64位x86程序")
        else:
            printer.test("PE文件运行平台: " + fMachine)
    else:
        printer.test("标志位错误,当前文件不是有效的PE文件!")

def pe_sections(path):
    printer.info("PE节分析: ")
    pe = pefile.PE(path)
    for section in pe.sections:
        sname = section.Name.decode().rstrip("\x00")
        if sname in peSection.keys():
            va = hex(section.VirtualAddress)
            vs = hex(section.Misc_VirtualSize)
            srd = hex(section.SizeOfRawData)
            praw = hex(section.PointerToRawData)
            printer.test(sname + " " + peSection[sname])
            printer.test("\n|\n|---- VirutalAddress[RVA] : " + hex(section.VirtualAddress) + " 节区的相对虚拟地址")
            printer.test("\n|\n|---- VirtualSize    : " + hex(section.Misc_VirtualSize) +" " +str(int(vs,16)) +" 节在内存中大小")
            printer.test("\n|\n|---- SizeOfRawData  : " + hex(section.SizeOfRawData) + " " +str(int(srd,16))+ " 节在文件中所占的大小")
            printer.test("\n|\n|---- PointerToRawData  : " + hex(section.PointerToRawData) + " " +str(int(praw,16))+ " 节在文件中的起始位置")
            printer.test("\n")

def pe_import_info(path):
    pe = pefile.PE(path)
    for entry in pe.DIRECTORY_ENTRY_IMPORT:
        printer.test("动态链接库: " + entry.dll.decode())
        for function in entry.imports:
            if function.name == None:
                continue
            printer.test("\t 导出的函数: " + function.name.decode())

def pe_dll_func(path,dll):
    pe = pefile.PE(path)
    for entry in pe.DIRECTORY_ENTRY_IMPORT:
        dll_name = entry.dll.decode('utf-8')
        if dll_name == dll:
            printer.test("动态链接库: " + dll_name)
            for function in entry.imports:
                if function.name == None:
                    continue
                printer.test("\t 调用的函数: " + function.name.decode() + " 内存地址: 0x%08x"%(function.address))
def pe_funcs(path):
    funcs = []
    pe = pefile.PE(path)
    for entry in pe.DIRECTORY_ENTRY_IMPORT:
        for function in entry.imports:
            if function.name == None:
                continue
            funcs.append(function.name.decode())
    return funcs

def pe_extract_strings(path):
    rfile = open(path, 'r',errors='ignore').read()
    strings = re.findall("[\x1f-\x7e]{4,}", rfile)
    return strings
def pe_extract_strings_range(strings,start,end):
    for s in strings[start:end]:
        print(s)
def pe_strings_find_mal(strings):
    for i in strings:
        for j in malstr.strings:
            if str(i).lower().find(str(j).lower()) >= 0:
                printer.warn("疑似字符串: " + i + ", 有效载荷: " + j)
    print("遍历完毕!")
def pe_func_find_mal(funcs:list):
    for f in funcs:
        for mf in malfunc.func.keys():
            if f.lower().find(mf.lower()) >= 0:
                printer.test("函数: " + f + " 描述: " + malfunc.func[mf])
def pe_get_hash(fp):
    pe = pefile.PE(fp)
    return pe.get_imphash()

#文件中的警告列表
def pe_get_warns(fp):
    pe = pefile.PE(fp)
    return pe.get_warnings()

#rva(内存偏移) to raw(文件偏移)
def pe_rva_raw(path,rva):
    pe = pefile.PE(path)
    #raw = rva - VirtualAddress + PointerToRawData
    raw = pe.get_offset_from_rva(rva)
    return raw
def pe_info(path,info="all"):
    pe = pefile.PE(path)
    if info == "dos":
        print(pe.DOS_HEADER)
    elif info == "nt":
        print(pe.NT_HEADERS)
        print(pe.FILE_HEADER)
        print(pe.OPTIONAL_HEADER)
    else:
        print(pe)


